import cv2
import numpy as np
import random
from gpeno.face_detect.utils.box_utils import matrix_iof


def _crop(image, boxes, labels, landm, img_dim):
	height, width, _ = image.shape
	pad_image_flag = True

	for _ in range(250):
		"""
        if random.uniform(0, 1) <= 0.2:
            scale = 1.0
        else:
            scale = random.uniform(0.3, 1.0)
        """
		PRE_SCALES = [0.3, 0.45, 0.6, 0.8, 1.0]
		scale = random.choice(PRE_SCALES)
		short_side = min(width, height)
		w = int(scale * short_side)
		h = w

		if width == w:
			l = 0
		else:
			l = random.randrange(width - w)
		if height == h:
			t = 0
		else:
			t = random.randrange(height - h)
		roi = np.array((l, t, l + w, t + h))

		value = matrix_iof(boxes, roi[np.newaxis])
		flag = (value >= 1)
		if not flag.any():
			continue

		centers = (boxes[:, :2] + boxes[:, 2:]) / 2
		mask_a = np.logical_and(roi[:2] < centers, centers < roi[2:]).all(axis=1)
		boxes_t = boxes[mask_a].copy()
		labels_t = labels[mask_a].copy()
		landms_t = landm[mask_a].copy()
		landms_t = landms_t.reshape([-1, 5, 2])

		if boxes_t.shape[0] == 0:
			continue

		image_t = image[roi[1]:roi[3], roi[0]:roi[2]]

		boxes_t[:, :2] = np.maximum(boxes_t[:, :2], roi[:2])
		boxes_t[:, :2] -= roi[:2]
		boxes_t[:, 2:] = np.minimum(boxes_t[:, 2:], roi[2:])
		boxes_t[:, 2:] -= roi[:2]

		# landm
		landms_t[:, :, :2] = landms_t[:, :, :2] - roi[:2]
		landms_t[:, :, :2] = np.maximum(landms_t[:, :, :2], np.array([0, 0]))
		landms_t[:, :, :2] = np.minimum(landms_t[:, :, :2], roi[2:] - roi[:2])
		landms_t = landms_t.reshape([-1, 10])

		# make sure that the cropped image contains at least one face > 16 pixel at training image scale
		b_w_t = (boxes_t[:, 2] - boxes_t[:, 0] + 1) / w * img_dim
		b_h_t = (boxes_t[:, 3] - boxes_t[:, 1] + 1) / h * img_dim
		mask_b = np.minimum(b_w_t, b_h_t) > 0.0
		boxes_t = boxes_t[mask_b]
		labels_t = labels_t[mask_b]
		landms_t = landms_t[mask_b]

		if boxes_t.shape[0] == 0:
			continue

		pad_image_flag = False

		return image_t, boxes_t, labels_t, landms_t, pad_image_flag
	return image, boxes, labels, landm, pad_image_flag


def _distort(image):

	def _convert(image, alpha=1, beta=0):
		tmp = image.astype(float) * alpha + beta
		tmp[tmp < 0] = 0
		tmp[tmp > 255] = 255
		image[:] = tmp

	image = image.copy()

	if random.randrange(2):

		#brightness distortion
		if random.randrange(2):
			_convert(image, beta=random.uniform(-32, 32))

		#contrast distortion
		if random.randrange(2):
			_convert(image, alpha=random.uniform(0.5, 1.5))

		image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)

		#saturation distortion
		if random.randrange(2):
			_convert(image[:, :, 1], alpha=random.uniform(0.5, 1.5))

		#hue distortion
		if random.randrange(2):
			tmp = image[:, :, 0].astype(int) + random.randint(-18, 18)
			tmp %= 180
			image[:, :, 0] = tmp

		image = cv2.cvtColor(image, cv2.COLOR_HSV2BGR)

	else:

		#brightness distortion
		if random.randrange(2):
			_convert(image, beta=random.uniform(-32, 32))

		image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)

		#saturation distortion
		if random.randrange(2):
			_convert(image[:, :, 1], alpha=random.uniform(0.5, 1.5))

		#hue distortion
		if random.randrange(2):
			tmp = image[:, :, 0].astype(int) + random.randint(-18, 18)
			tmp %= 180
			image[:, :, 0] = tmp

		image = cv2.cvtColor(image, cv2.COLOR_HSV2BGR)

		#contrast distortion
		if random.randrange(2):
			_convert(image, alpha=random.uniform(0.5, 1.5))

	return image


def _expand(image, boxes, fill, p):
	if random.randrange(2):
		return image, boxes

	height, width, depth = image.shape

	scale = random.uniform(1, p)
	w = int(scale * width)
	h = int(scale * height)

	left = random.randint(0, w - width)
	top = random.randint(0, h - height)

	boxes_t = boxes.copy()
	boxes_t[:, :2] += (left, top)
	boxes_t[:, 2:] += (left, top)
	expand_image = np.empty((h, w, depth), dtype=image.dtype)
	expand_image[:, :] = fill
	expand_image[top:top + height, left:left + width] = image
	image = expand_image

	return image, boxes_t


def _mirror(image, boxes, landms):
	_, width, _ = image.shape
	if random.randrange(2):
		image = image[:, ::-1]
		boxes = boxes.copy()
		boxes[:, 0::2] = width - boxes[:, 2::-2]

		# landm
		landms = landms.copy()
		landms = landms.reshape([-1, 5, 2])
		landms[:, :, 0] = width - landms[:, :, 0]
		tmp = landms[:, 1, :].copy()
		landms[:, 1, :] = landms[:, 0, :]
		landms[:, 0, :] = tmp
		tmp1 = landms[:, 4, :].copy()
		landms[:, 4, :] = landms[:, 3, :]
		landms[:, 3, :] = tmp1
		landms = landms.reshape([-1, 10])

	return image, boxes, landms


def _pad_to_square(image, rgb_mean, pad_image_flag):
	if not pad_image_flag:
		return image
	height, width, _ = image.shape
	long_side = max(width, height)
	image_t = np.empty((long_side, long_side, 3), dtype=image.dtype)
	image_t[:, :] = rgb_mean
	image_t[0:0 + height, 0:0 + width] = image
	return image_t


def _resize_subtract_mean(image, insize, rgb_mean):
	interp_methods = [cv2.INTER_LINEAR, cv2.INTER_CUBIC, cv2.INTER_AREA, cv2.INTER_NEAREST, cv2.INTER_LANCZOS4]
	interp_method = interp_methods[random.randrange(5)]
	image = cv2.resize(image, (insize, insize), interpolation=interp_method)
	image = image.astype(np.float32)
	image -= rgb_mean
	return image.transpose(2, 0, 1)


class preproc(object):

	def __init__(self, img_dim, rgb_means):
		self.img_dim = img_dim
		self.rgb_means = rgb_means

	def __call__(self, image, targets):
		assert targets.shape[0] > 0, "this image does not have gt"

		boxes = targets[:, :4].copy()
		labels = targets[:, -1].copy()
		landm = targets[:, 4:-1].copy()

		image_t, boxes_t, labels_t, landm_t, pad_image_flag = _crop(image, boxes, labels, landm, self.img_dim)
		image_t = _distort(image_t)
		image_t = _pad_to_square(image_t, self.rgb_means, pad_image_flag)
		image_t, boxes_t, landm_t = _mirror(image_t, boxes_t, landm_t)
		height, width, _ = image_t.shape
		image_t = _resize_subtract_mean(image_t, self.img_dim, self.rgb_means)
		boxes_t[:, 0::2] /= width
		boxes_t[:, 1::2] /= height

		landm_t[:, 0::2] /= width
		landm_t[:, 1::2] /= height

		labels_t = np.expand_dims(labels_t, 1)
		targets_t = np.hstack((boxes_t, landm_t, labels_t))

		return image_t, targets_t
